package cas.ihep.spark.rdd;

import cas.ihep.spark.context.JSparkIterationListener;
import cas.ihep.spark.loop.ClientIterativeRunnable;
import cas.ihep.spark.loop.IterationClient;
import cas.ihep.spark.loop.IterationData;
import cas.ihep.spark.loop.IterativeRunnable;
import cas.ihep.spark.util.GlobalConfiguration;
import cas.ihep.spark.util.IterationClientListener;
import cas.ihep.spark.util.NamedResourceManager;
import cas.ihep.spark.util.function.Function2;
import org.apache.spark.SparkContext;
import org.apache.spark.api.java.function.VoidFunction;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.rdd.RDD;
import scala.reflect.ClassManifestFactory;
import scala.runtime.AbstractFunction1;
import scala.runtime.BoxedUnit;

import java.io.Serializable;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.Map;

final class ClientIterativeRunnableImpl<T,S extends Serializable,R extends Serializable> implements ClientIterativeRunnable<S,R>{

    private T baseValue;
    private NamedResourceManager nm;
    private IterativeBroadcastDatav2<T,S,R> bcastData;

    ClientIterativeRunnableImpl(T b,NamedResourceManager n,IterativeBroadcastDatav2<T,S,R> bd){
        baseValue=b;
        nm=n;
        bcastData=bd;
    }
    @Override
    public void run(IterationData<S, R> clientData) {
        clientData.output(bcastData.clientListener.iterate(nm,baseValue,clientData.input()));
    }
}

abstract class IterativeRDDUtil<T> extends AbstractFunction1<T, BoxedUnit> implements VoidFunction<T> {
    static <T,Send extends Serializable, Receive extends Serializable>
    void iterate(RDD<T> rdd, SparkContext sc, int numPartitions, IterativeRunnable<Send, Receive> iRun, IterationClientListener<T,Send,Receive> cRun,
                 Class<Send> sClass, Receive zero, Function2<Receive,Receive,Receive> aggregateFunc){
        IterativeBroadcastDatav2<T, Send, Receive> param;
        try {
            param=new IterativeBroadcastDatav2<>(new InetSocketAddress(InetAddress.getLocalHost(), 10000),sClass,
                    zero,aggregateFunc,cRun);
        }catch (UnknownHostException error){
            throw new RuntimeException("Should never happens",error);
        }
        param.properties= GlobalConfiguration.properties();
        sc.addSparkListener(new JSparkIterationListener<>(sc, iRun, numPartitions, param.serverAddr));
        final Broadcast<IterativeBroadcastDatav2<T, Send, Receive>> bcast=sc.broadcast(param, ClassManifestFactory.classType(param.getClass()));
        rdd.foreach(new IterativeRDDUtil<T>() {
            @Override
            public void call(T t) throws Exception {
                IterativeBroadcastDatav2<T, Send, Receive> param1=bcast.getValue();
                initExecutorProperties(param1.properties);
                try(NamedResourceManagerImpl closer=new NamedResourceManagerImpl()){
                    InetSocketAddress localAddr=new InetSocketAddress(InetAddress.getLocalHost(),8999);
                    @SuppressWarnings("unchecked")
                    IterationClient<Send, Receive> client1=IterationClient.getOrCreateSingleInstance(BaseBroadcastData.CLIENT_KEY,localAddr,
                            param1.serverAddr,(Class<Send>)param1.sendClassTag.runtimeClass());
                    param1.clientListener.initialize(closer,t);
                    client1.waitingForData(new ClientIterativeRunnableImpl<>(t,closer,param1),param1.zero,param1.aggregateFunc);
                    param1.clientListener.cleanup(closer,t);
                }
            }
        });
    }

    static void initExecutorProperties(Map<String,String> props){
        if(!inited){
            synchronized (IterativeRDDUtil.class){
                if(!inited){
                    GlobalConfiguration.update(props);
                    inited=true;
                }
            }
        }
    }

    private static volatile boolean inited=false;

    public abstract void call(T t)throws Exception;

    public final BoxedUnit apply(T t){
        try {
            call(t);
        }catch (RuntimeException e){
            throw e;
        }catch (Exception e){
            throw new RuntimeException(e);
        }
        return BoxedUnit.UNIT;
    }
}
